<!DOCTYPE html>
<html>
<head>
  <title>tinygrad viz</title>
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <link rel="icon" href="data:;base64,iVBORw0KGgo=">
  <script src="assets/d3js.org/d3.v5.min.js" charset="utf-8"></script>
  <script src="assets/dagrejs.github.io/project/dagre-d3/latest/dagre-d3.min.js"></script>
  <link rel="stylesheet" href="assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/styles/default.min.css">
  <script src="assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/highlight.min.js"></script>
  <script src="assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/python.min.js"></script>
  <script src="assets/cdnjs.cloudflare.com/ajax/libs/highlight.js/11.10.0/languages/cpp.min.js"></script>
  <script src="assets/cdnjs.cloudflare.com/ajax/libs/dompurify/1.0.3/purify.min.js"></script>
  <link rel="stylesheet" href="assets/unpkg.com/@highlightjs/cdn-assets@11.10.0/styles/tokyo-night-dark.min.css" />
  <style>
  * {
    box-sizing: border-box;
    margin-block-start: initial;
    margin-block-end: initial;
  }
  button {
    outline: none;
  }
  html, body {
    color: #f0f0f5;
    margin: 0;
    padding: 0;
    width: 100%;
    height: 100%;
    font-family: "Noto Sans", sans-serif;
    font-optical-sizing: auto;
    font-weight: 400;
    font-style: normal;
    font-variation-settings: "wdth" 100;
    font-size: 14px;
    overflow: hidden;
  }
  a {
    color: #4a90e2;
  }
  ul {
    padding: 0;
    color: #7c7d85;
    white-space: nowrap;
    cursor: pointer;
  }
  ul.active {
    color: #f0f0f5;
  }
  svg {
    width: 100%;
    height: 100%;
  }
  svg * {
    cursor: default;
    user-select: none;
  }
  .node rect {
    stroke: #4a4b57;
    stroke-width: 1.4px;
    rx: 8px;
    ry: 8px;
  }
  .label text {
    color: #08090e;
    font-weight: 350;
  }
  .edgePath path {
    stroke: #4a4b57;
    fill: #4a4b57;
    stroke-width: 1.4px;
  }
  .main-container {
    display: flex;
    width: 100%;
    height: 100%;
    position: relative;
  }
  .container {
    background-color: #0f1018;
  }
  .container > * + *, .rewrite-container > * + * {
    margin-top: 12px;
  }
  .graph {
    background-color: #08090e;
    position: absolute;
    inset: 0;
    z-index: 1;
  }
  .kernel-list-parent {
    position: relative;
    width: 15%;
    padding: 50px 20px 20px 20px;
    border-right: 1px solid #4a4b56;
    z-index: 2;
  }
  .kernel-list {
    width: 100%;
    height: 100%;
    overflow-y: auto;
  }
  .kernel-list > ul > * + * {
    margin-top: 4px;
  }
  .metadata {
    position: relative;
    width: 20%;
    padding: 20px;
    background-color: #0f1018;
    border-left: 1px solid #4a4b56;
    z-index: 2;
    margin-left: auto;
    height: 100%;
    overflow-y: auto;
  }
  .resize-handle {
    position: absolute;
    top: 0;
    bottom: 0;
    width: 20px;
    height: 100%;
    cursor: col-resize;
    z-index: 3;
    background-color: transparent;
  }
  #kernel-resize-handle {
    right: 0;
  }
  #metadata-resize-handle {
    left: 0;
  }
  .floating-container {
    position: fixed;
    top: 10px;
    left: 20px;
    z-index: 4;
    display: flex;
    flex-direction: row;
    gap: 8px;
  }
  .nav-btn {
    background-color: #1a1b26;
    border: 1px solid #4a4b56;
    color: #f0f0f5;
    height: 32px;
    border-radius: 8px;
    cursor: pointer;
    text-decoration: none;
    display: flex;
    align-items: center;
    padding: 0 6px;
    font-weight: bold;
  }
  .collapse-btn {
    width: 32px;
    padding: 6px;
  }
  .btn {
    height: 32px;
    background-color: #1a1b26;
    border: 1px solid #4a4b56;
    color: #f0f0f5;
    border-radius: 8px;
    cursor: pointer;
    transition-duration: .5s;
  }
  .btn:hover {
    background-color: #2a2b36;
    border-color: #5a5b66;
    color: #ffffff;
    transform: translateY(-1px);
  }
  .collapsed .kernel-list, .collapsed .metadata {
    width: 0;
    padding: 0;
    overflow: hidden;
  }
  .rewrite-list {
    display: flex;
    flex-wrap: wrap;
  }
  .rewrite-list > * + * {
    margin-left: 4px;
  }
  .wrap {
    word-wrap: break-word;
    white-space: pre-wrap;
  }
  .code-block {
    max-height: 30%;
    overflow-y: auto;
    border-radius: 8px;
    padding: 8px;
  }
  </style>
</head>
<body>
  <div class="main-container">
    <div class="floating-container">
      <button class="btn collapse-btn">
        <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M15 19l-7-7 7-7"/></svg>
      </button>
      <a class="btn nav-btn" href="/profiler">Profiler</a>
    </div>
    <div class="container kernel-list-parent"><div class="container kernel-list"></div></div>
    <div class="graph">
      <svg id="graph-svg">
        <g id="render"></g>
      </svg>
    </div>
    <div class="container metadata"></div>
  </div>
  <script>
  /* global hljs, dagreD3, d3, DOMPurify */

  // **** hljs extra definitions for UOps and float4
  hljs.registerLanguage("python", (hljs) => ({
    ...hljs.getLanguage("python"),
    case_insensitive: false,
    contains: [
      { begin: 'dtypes\\.[a-zA-Z_][a-zA-Z0-9_-]*(\\.[a-zA-Z_][a-zA-Z0-9_-]*)*' + '(?=[.\\s\\n[:,(])', className: "type" },
      { begin: 'dtypes\\.[a-zA-Z_][a-zA-Z0-9_-].vec*' + '(?=[.\\s\\n[:,(])', className: "type" },
      { begin: '[a-zA-Z_][a-zA-Z0-9_-]*\\.[a-zA-Z_][a-zA-Z0-9_-]*' + '(?=[.\\s\\n[:,()])',  className: "operator" },
      { begin: '[A-Z][a-zA-Z0-9_]*(?=\\()', className: "section", ignoreEnd: true },
      ...hljs.getLanguage("python").contains,
    ]
  }));
  hljs.registerLanguage("cpp", (hljs) => ({
    ...hljs.getLanguage('cpp'),
    contains: [{ begin: '\\b(?:float|half)[0-9]+\\b', className: 'type' }, ...hljs.getLanguage('cpp').contains]
  }));

  // **** D3
  function recenterRects(svg, zoom) {
    const svgBounds = svg.node().getBoundingClientRect();
    for (const rect of svg.node().querySelectorAll("rect")) {
      const rectBounds = rect.getBoundingClientRect();
      const outOfBounds = rectBounds.top < svgBounds.top || rectBounds.left < svgBounds.left ||
        rectBounds.bottom > svgBounds.bottom || rectBounds.right > svgBounds.right;
      // if there's at least one rect in view we don't do anything
      if (!outOfBounds) return;
    }
    svg.call(zoom.transform, d3.zoomIdentity)
  }
  function renderGraph(graph, additions) {
    const g = new dagreD3.graphlib.Graph({ compound: true }).setGraph({ rankdir: "LR" }).setDefaultEdgeLabel(function() { return {}; });
    g.setNode("addition", {label: "", clusterLabelPos: "top", style: additions.length !== 0 ? "fill: rgba(26, 27, 38, 0.5);" : "display: none;"});
    for (const [k,u] of Object.entries(graph)) {
      g.setNode(k, {label: u[0], style: `fill: ${u[4]};` });
      for (const src of u[2]) {
        g.setEdge(src, k, {curve: d3.curveBasis})
      }
      if (additions.includes(parseInt(k))) {
        g.setParent(k, "addition");
      }
    }
    const svg = d3.select("#graph-svg");
    const inner = svg.select("g");
    var zoom = d3.zoom()
      .scaleExtent([0.05, 2])
      .on("zoom", () => {
        const transform = d3.event.transform;
        inner.attr("transform", transform);
      });
    recenterRects(svg, zoom);
    svg.call(zoom);
    const render = new dagreD3.render();
    render(inner, g);
  }

  // **** extra helpers
  const toPath = ([fp, lineno]) => `${fp.replaceAll("\\", "/").split("/").pop()}:${lineno}`;
  const vsCodeOpener = (parts) => Object.assign(document.createElement("a"), { textContent: parts[parts.length-1]+"\n\n",
      href: "vscode://file"+parts.join("/"), style: "font-family: monospace; margin: 4px 0;" });

  // **** main loop
  var ret = {};
  var cache = {};
  var kernels = null;
  var currentUOp = 0;
  var currentKernel = -1;
  var currentRewrite = 0;
  var expandKernel = true;
  async function main() {
    const mainContainer = document.querySelector('.main-container');
    // ***** LHS kernels list
    if (kernels == null) {
      kernels = await (await fetch("/kernels")).json();
      currentKernel = -1;
    }
    const kernelListParent = document.querySelector(".container.kernel-list-parent");
    const kernelList = document.querySelector(".container.kernel-list");
    kernelList.innerHTML = "";
    kernels.forEach((k, i) => {
      const kernelUl = Object.assign(document.createElement("ul"), { key: `kernel-${i}`, className: i === currentKernel ? "active" : "",
        style: "overflow-x: auto; cursor: initial;" });
      if (i === currentKernel) {
        requestAnimationFrame(() => kernelUl.scrollIntoView({ behavior: "auto", block: "nearest" }));
      }
      const p = Object.assign(document.createElement("p"), { id: `kernel-${k[0].kernel_name}`, innerText: k[0].kernel_name ?? "UNPARENTED",
        style: "cursor: pointer;"});
      kernelUl.appendChild(p)
      k.forEach((u, j) => {
        const rwUl = Object.assign(document.createElement("ul"), { innerText: `${toPath(u.loc)} - ${u.upats.length}`, key: `uop-rewrite-${j}`,
          className: (j === currentUOp && i == currentKernel) ? "active" : "" })
        if (j === currentUOp) {
          requestAnimationFrame(() => rwUl.scrollIntoView({ behavior: "auto", block: "nearest" }));
        }
        rwUl.style.display = i === currentKernel && expandKernel ? "block" : "none";
        rwUl.onclick = (e) => {
          e.stopPropagation();
          currentUOp = j;
          currentKernel = i;
          currentRewrite = 0;
          main();
        }
        kernelUl.appendChild(rwUl)
      })
      p.onclick = () => {
        if (i === currentKernel) {
          expandKernel = !expandKernel;
          main();
          return;
        }
        currentKernel = i;
        currentUOp = 0;
        currentRewrite = 0;
        expandKernel = true;
        main();
      }
      kernelList.appendChild(kernelUl);
    });
    // ***** UOp graph
    if (currentKernel != -1) {
      const cacheKey = `${currentKernel}-${currentUOp}`;
      if (cacheKey in cache) {
        ret = cache[cacheKey];
      }
      else {
        ret = await (await fetch(`/kernels?kernel=${currentKernel}&idx=${currentUOp}`)).json();
        cache[cacheKey] = ret;
      }
      renderGraph(ret.graphs[currentRewrite], currentRewrite == 0 ? [] : ret.changed_nodes[currentRewrite-1]);
    }
    // ***** RHS metadata
    const metadata = document.querySelector(".container.metadata");
    metadata.innerHTML = "";
    metadata.appendChild(vsCodeOpener(ret.loc.join(":").split("/")));
    const pre = Object.assign(document.createElement("pre"), { innerHTML: `<code>${DOMPurify.sanitize(ret.code_line)}</code>`,
      className: "wrap code-block language-python" });
    metadata.appendChild(pre);
    hljs.highlightElement(pre);
    // ** code blocks
    let code = ret.uops[currentRewrite];
    let lang = "python"
    if (ret.kernel_code != null) {
      code = ret.kernel_code.replaceAll("<", "&lt;").replaceAll(">", "&gt;");
      lang = "cpp";
    }
    const codeBlock = Object.assign(document.createElement("pre"), { innerHTML: `<code>${DOMPurify.sanitize(code)}</code>`,
      className: `code-block language-${lang}` });
    hljs.highlightElement(codeBlock);
    metadata.appendChild(codeBlock);
    // ** rewrite list
    if (ret.graphs.length > 1) {
      const rewriteList = Object.assign(document.createElement("div"), { className: "rewrite-list" })
      metadata.appendChild(rewriteList);
      ret.graphs.forEach((rw, i) => {
        const gUl = Object.assign(document.createElement("ul"), { innerText: i, key: `rewrite-${i}` });
        rewriteList.appendChild(gUl);
        if (i === currentRewrite) {
          gUl.classList.add("active");
          if (i !== 0) {
            const diff = ret.diffs[i-1];
            const [loc, pattern] = ret.upats[i-1];
            const parts = loc.join(":").split("/");
            const div = Object.assign(document.createElement("div"), { className: "rewrite-container" });
            const link = vsCodeOpener(parts);
            div.appendChild(link);
            const pre = Object.assign(document.createElement("pre"), { className: "code-block wrap" });
            pre.appendChild(Object.assign(document.createElement("code"), { textContent: DOMPurify.sanitize(pattern), className: "language-python" }));
            div.appendChild(pre);
            hljs.highlightElement(pre);
            metadata.appendChild(div);
            const diffHtml = diff.map((line) => {
              const color = line.startsWith("+") ? "#3aa56d" : line.startsWith("-") ? "#d14b4b" : "#f0f0f5";
              return `<span style="color: ${color};">${line}</span>`;
            }).join("<br>");
            metadata.appendChild(Object.assign(document.createElement("pre"), { innerHTML: `<code>${diffHtml}</code>`, className: "wrap" }));
          }
        }
        gUl.addEventListener("click", () => {
          currentRewrite = i;
          main();
        });
      })
    } else {
      metadata.appendChild(Object.assign(document.createElement("p"), { textContent: `No rewrites in ${toPath(ret.loc)}.` }));
    }
    // ***** collapse/expand
    let isCollapsed = false;
    const collapseBtn = document.querySelector(".collapse-btn");
    collapseBtn.addEventListener("click", () => {
      isCollapsed = !isCollapsed;
      mainContainer.classList.toggle("collapsed", isCollapsed);
      kernelListParent.style.display = isCollapsed ? "none" : "block";
      metadata.style.display = isCollapsed ? "none" : "block";
      collapseBtn.style.transform = isCollapsed ? "rotate(180deg)" : "rotate(0deg)";
    });
    // ***** resizer
    function createResizer(element, width, type) {
      const { minWidth, maxWidth } = width;
      const handle = Object.assign(document.createElement("div"), { id: `${type}-resize-handle`, className: "resize-handle" });
      element.appendChild(handle);

      const resize = (e) => {
        const change = e.clientX - element.dataset.startX;
        const adjustedChange = type === "kernel" ? change : -change;
        const newWidth = ((Number(element.dataset.startWidth) + adjustedChange) / Number(element.dataset.containerWidth)) * 100;
        if (newWidth >= minWidth && newWidth <= maxWidth) {
          element.style.width = `${newWidth}%`;
        }
      };

      handle.addEventListener("mousedown", (e) => {
        e.preventDefault();
        element.dataset.startX = e.clientX;
        element.dataset.containerWidth = mainContainer.getBoundingClientRect().width;
        element.dataset.startWidth = element.getBoundingClientRect().width;

        document.documentElement.addEventListener("mousemove", resize, false);
        document.documentElement.addEventListener("mouseup", () => {
          document.documentElement.removeEventListener("mousemove", resize, false);
          element.style.userSelect = "initial";
        }, { once: true });
      });
    }
    createResizer(kernelListParent, { minWidth: 15, maxWidth: 50 }, "kernel"); // left resizer
    createResizer(metadata, { minWidth: 20, maxWidth: 50 }, "metadata"); // right resizer
  }

  // **** keyboard shortcuts
  document.addEventListener("keydown", async function(event) {
    // up and down change the UOp or kernel from the list
    if (!expandKernel) {
      if (event.key == "ArrowUp") {
        event.preventDefault()
        currentUOp = 0;
        currentRewrite = 0;
        currentKernel = Math.max(0, currentKernel-1)
       return main()
      }
      if (event.key == "ArrowDown") {
        event.preventDefault()
        currentUOp = 0;
        currentRewrite = 0;
        currentKernel = Math.min(Array.from(Object.keys(kernels)).length-1, currentKernel+1)
       return main()
      }
    }
    if (event.key == "Enter") {
      event.preventDefault()
      if (currentKernel === -1) {
        currentKernel = 0;
        expandKernel = true;
      }
      else {
        expandKernel = !expandKernel;
      }
      currentUOp = 0;
      currentRewrite = 0;
      main();
    }
    if (event.key == "ArrowUp") {
      event.preventDefault()
      currentRewrite = 0;
      currentUOp = Math.max(0, currentUOp-1)
      main()
    }
    if (event.key == "ArrowDown") {
      event.preventDefault()
      currentRewrite = 0;
      const totalUOps = kernels[currentKernel].length-1;
      currentUOp = Math.min(totalUOps, currentUOp+1)
      main()
    }
    // left and right go through rewrites in a single UOp
    if (event.key == "ArrowLeft") {
      event.preventDefault()
      currentRewrite = Math.max(0, currentRewrite-1)
      main()
    }
    if (event.key == "ArrowRight") {
      event.preventDefault()
      const totalRewrites = ret.graphs.length-1;
      currentRewrite = Math.min(totalRewrites, currentRewrite+1)
      main()
    }
  })
  main()
  </script>
  </body>
</html>
